#include <iostream>
#include <vector>
#include <ctime>
#include <numeric>

#include "tensor.h"
#include "covariance.h"
#include "mtrand.h"
#include "dataset.h"
#include "pdf.h"
#include <vector>
#include <sstream>
#include <fstream>

#include "TFile.h"
#include "TChain.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TProfile2D.h"
#include "TF1.h"
#include "TF2.h"
#include "TF3.h"
#include "TLorentzVector.h"
#include "TGraphErrors.h"
#include "TGraph2DErrors.h"
#include "make_list.h"

using namespace std;

//initialize random number generator
MTRand_int32 irand(time(0));

const double PI=3.14159265359;
double min_pt=150.;
const bool do_MC=false;

//counter for number of events
class Counter{
  
public:

  double n;
  double n2;
  double error;
  double tot_n2;
  double tot_n;
  
  
  Counter()
  {
    n=0;
    n2=0;
    error=0;
    tot_n=0;
    tot_n2=0;
  }

  void Add(double weight, double error_=0)
  {
    n+=weight;
    n2+=weight*weight;
    error+=error_;
  }

  
  void Add_lhc(double evt_weight, double frac, double error_=0)
  {
    n+=evt_weight*frac;
    n2+=evt_weight*frac*frac;
    error+=error_;
  }
  
  void Add_tot_n(double tot_err)
  {
    tot_n += tot_err;
    tot_n2 += tot_err*tot_err;
  }

  void Add_tot_n_lhc(double evt_weight, double tot_err)
  {
    tot_n += evt_weight*tot_err;
    tot_n2 += evt_weight*tot_err*tot_err;
  }

  double tot_err() const
  {
    return n*sqrt(tot_n2)/tot_n;
  }

  double nevt() const
  {
    return n;
  }
    
  double err() const
  {
    return sqrt(n2 + error*error);
  }

  double bias() const
  {
    return error;
  }

  double var() const
  {
    return sqrt(n2);
  }

};




//function to produce a 2D plot from DataSet
TGraph2D plot2D(DataSet& data, char* name)
{
  //ensure that dataset is two dimensional
  assert(data.Rank() == 2);

  PDF pdf=data.ComputePDF();
  
  //declare some x,y points
  vector<double> x;
  vector<double> y;
  vector<double> val;
  
  const int nbin=20;

  double x_diff=(data.max[0]-data.min[0])/nbin;
  double y_diff=(data.max[1]-data.min[1])/nbin;

  for(unsigned int i=0; i<nbin; ++i)
  for(unsigned int j=0; j<nbin; ++j)
    {
      double x_val=data.min[0]+x_diff*(i+0.5);
      double y_val=data.min[1]+y_diff*(j+0.5);

      x.push_back(x_val);
      y.push_back(y_val);
      
      val.push_back(pdf(make_list<double>(x_val, y_val)));
    }
    
  //get a new TGraph2D
  TGraph2D graph(x.size(), &x[0], &y[0], &val[0]);
  graph.SetName(name);
  graph.SetTitle(name);

  return graph;
}


class Vec
{
public:
  vector<double> x,y,err;

  Vec(){}

  void push_back(double x_, double y_, double yerr_)
  {
    x.push_back(x_);
    y.push_back(y_);
    err.push_back(yerr_);    
  }

  TGraphErrors get_graph(string name, bool logx=false) const
  {
    vector<double> xprime(x);
    if(logx)
      {
	for(unsigned int i=0; i<xprime.size(); ++i)
	  {
	    xprime[i]=exp10(-2*xprime[i]);
	  }
      }
    TGraphErrors result(xprime.size(), &xprime[0], &y[0], NULL, &err[0]);
    result.SetTitle(name.c_str());
    result.SetName(name.c_str());

    return result;
  }

  //get a 1D graph from pt
  void fill_data(PDF& pdf, double pt)
  {
    PDF::MC_Generator gen=
      pdf.Generator(make_list<PDF::MC_Flag>
		    (PDF::OUTPUT, PDF::INPUT));

    vector<PDF::MC_Event> evt=
      gen.GenerateFull(make_list<double>(log(pt/min_pt)));

    for(unsigned int i=0; i<evt.size(); ++i)
      {
	x.push_back(evt[i].value[0]);
	y.push_back(evt[i].weight);
	err.push_back(evt[i].err);
      }
  }

};



float DeltaR(float eta, float phi)
{
  if(fabs(phi)> PI)
    phi= PI-fabs(phi);
  return sqrt(eta*eta + phi*phi);
}

int main(int argc, char **argv)
{
  if(argc < 2+1)
    {
      cout<<"Usage: ./madgraph_multijet_cuttable madgraph_input output.root [extra smoothing factor] [random seed if PDF_smearing on] [lumi] [min_pt]"
	  <<endl;
      throw;
    }

  cout<<"madgraph inputs are read as madgraph_input.?.root ranges from 1 to 12"
      <<endl;
  cout<<"Outputs are automatically stored in PDF_*"<<endl;
  
  //see if we want to fluctuate the PDF to get variance
  bool b_smear_PDF=false;
  bool b_smear_lhc=false;

  //scale up to 20 fb^{-1}
  //weights are in milibarns, hence scale up by 1e12 to femtobarn
  double lhc_factor= 100.0* 1e9;

  if(argc >= 3+1)
    {
      cout<<"Extra smoothing factor on "<<atof(argv[3])<<endl;
      DataSet::kernel_smoothing_factor=atof(argv[3]);
    }



  if(argc >= 4+1)
    {
      b_smear_PDF=true;

      //see if there are extra arguments
      if(argc >= 5+1 && string(argv[5])=="LHC")
	{
	  b_smear_lhc=true;
	  cout<<"LHC mode enabled"<<endl;
	}

      if(argc >= 6+1 && atof(argv[6]) > 0)
	{
	  //lumi in picobarn
	  lhc_factor= atof(argv[6])*1e9;

	  cout<<"scale up to "<<atof(argv[6])<<" inverse picobarn"<<endl;


	  if(argc>= 7+1)
	    {
	      cout<<"change min_pt: "<<atoi(argv[7])<<endl;
	      min_pt=floor(atoi(argv[7]));
	    }
	  
	}
      
      cout<<"PDF smearing on, with seed: "<<atoi(argv[4])<<endl;

      //set random seed
      DataSet::SetRandomSeed(atoi(argv[4]));
    }


  //declare variables
  const int maxsize=20;
  Int_t           size;
  Double_t         w;   //[size]
  Float_t         pt[maxsize];   //[size]
  Float_t         eta[maxsize];
  Float_t         phi[maxsize];
  Float_t         m[maxsize];   //[size]
  Float_t         t1[maxsize];
  Float_t         t2[maxsize];
  Float_t         t3[maxsize];

  //exclusive 3jet
  //this is the number of jets for the training sample
  int n_exc=3;
  
  //list of DataSet
  vector<DataSet*> dataset_mpt_njet(n_exc);
  vector<DataSet*> dataset_tau21_njet(n_exc);
  vector<DataSet*> dataset_mpt_tau21_njet(n_exc);

  //Statistics plots for different Pt
  TH2D stat_template_1("stat_template_1", "Template Stats (50, 136)",
		       100, 0, 3, 100, 0, 1.0);

  TH2D stat_template_2("stat_template_2", "Template Stats (136, 370)",
		       100, 0, 3, 100, 0, 1.0);

  TH2D stat_template_3("stat_template_3", "Template Stats (370, 1004)",
		       100, 0, 3, 100, 0, 1.0);

  TH2D stat_template_4("stat_template_4", "Template Stats (1004, 2730)",
		       100, 0, 3, 100, 0, 1.0);

  TH2D stat_template_5("stat_template_5", "Template Stats (2730, 7421)",
		       100, 0, 3, 100, 0, 1.0);
  
  
  TH2D* lhc_stat[5]={&stat_template_1, &stat_template_2, &stat_template_3,
		     &stat_template_4, &stat_template_5};

  vector<Counter> lhc_prescale(3, Counter());
  double lhc_prescale_factor[3];
  double lhc_prescale_cut[3]={min_pt, floor(sqrt(280*min_pt)/5)*5, 280};
  
  cout<<"prescale bracket: "<<lhc_prescale_cut[0]
      <<", "<<lhc_prescale_cut[1]
      <<", "<<lhc_prescale_cut[2]<<endl;

  const double lhc_Hz=100.;

  int lhc_size=sizeof(lhc_stat)/sizeof(TH2D*);
  
  for(int is=0; is<lhc_size; ++is)
    {
      lhc_stat[is]->Sumw2();
    }


  //limit to 2 jets
  for(int n=0; n<n_exc && n<2; n++)
    {
      dataset_mpt_njet[n]=new
	DataSet(make_list<double>(0, 0),
		make_list<double>(3.0, 2.8),
		make_list<double>(0.02, 0.05));
  
      dataset_tau21_njet[n]=new
	DataSet(make_list<double>(0, 0),
		make_list<double>(1.0, 2.8),
		make_list<double>(0.02, 0.05));

      if(n==1)
	dataset_mpt_tau21_njet[n]=dataset_mpt_tau21_njet[n-1];

      else
	dataset_mpt_tau21_njet[n]=new
	  DataSet(make_list<double>(0, 0, 0),
		  make_list<double>(2.5, 0.95, 2.8),
		  make_list<double>(0.007, 0.006, 0.007));
    }

  //now process madgraph events
  char alist[3]={'A', 'B', 'C'};

  //output the error contribution for each templates
  vector<vector<double> > error_contrib_50;
  vector<vector<double> > nevt_contrib_50;
  vector<vector<double> > error_contrib_100;
  vector<vector<double> > nevt_contrib_100;
  vector<vector<double> > error_contrib_500;
  vector<vector<double> > nevt_contrib_500;

  for(unsigned int k=1; k<=12; k++)
    {
      error_contrib_50.push_back(vector<double>() );
      error_contrib_100.push_back(vector<double>() );
      error_contrib_500.push_back(vector<double>() );

      nevt_contrib_50.push_back(vector<double>() );
      nevt_contrib_100.push_back(vector<double>() );
      nevt_contrib_500.push_back(vector<double>() );

      for(unsigned int ka=0; ka<3; ka++)
	{
	  error_contrib_50[k-1].push_back(0);
	  error_contrib_100[k-1].push_back(0);
	  error_contrib_500[k-1].push_back(0);

	  nevt_contrib_50[k-1].push_back(0);
	  nevt_contrib_100[k-1].push_back(0);
	  nevt_contrib_500[k-1].push_back(0);
	}
    }

  
  //keep the weights
  double weights[12][3];

  //first loop over everything to compute prescale factor
  for(unsigned int k=1; k<=12; k++)
    for(unsigned int ka=0; ka<3; ka++)
      {
	TChain chain_madgraph("Ntuple");
	stringstream sout;
	sout<<argv[1]<<alist[ka]<<"."<<k<<".root";
	chain_madgraph.Add(sout.str().c_str());
	
	//read the input file
	chain_madgraph.SetBranchAddress("size", &size);
	chain_madgraph.SetBranchAddress("w", &w);
	chain_madgraph.SetBranchAddress("pt", pt);
	chain_madgraph.SetBranchAddress("eta", eta);
	chain_madgraph.SetBranchAddress("phi", phi);
	chain_madgraph.SetBranchAddress("m", m);
	chain_madgraph.SetBranchAddress("tau1", t1);
	chain_madgraph.SetBranchAddress("tau2", t2);
	chain_madgraph.SetBranchAddress("tau3", t3);
	
	double rescale=1.0/chain_madgraph.GetEntries();
	
	for(unsigned int i=0; i<chain_madgraph.GetEntries(); ++i)
	  {
	    if(i%10000==0)
	      cout<<"\r\033[K"<<floor((i/(chain_madgraph.GetEntries()*12.*3)  +
				       (3*(k-1)+ka)/(12.*3)
				       )*10000)*0.01
		  <<"% pre-processed (for prescale)..."
		  <<flush;
	  

	    chain_madgraph.GetEntry(i);
	    if(size<2) continue;

	    w*=rescale;

	    weights[k-1][ka]+=w;

	    
	    int real_njet=0;
	    for(int n=0; n<size; ++n)
	      {
		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;
		if(pt[n] < min_pt) continue;
		if(m[n]<20) continue;
		real_njet++;
	      }
	    
	    //if(real_njet!=n_exc) continue;
	    
	    int n_real=0;
	    //limit to two jets
	    for(int n=0; n<size && n_real<n_exc && n_real<2; ++n)
	      {
		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;
		if(pt[n] < min_pt) continue;
		if(m[n]<20) continue;
		
		/*
		double logpt=log(pt[n]/min_pt);
		double mpt=-log10(m[n]/pt[n]);
		double t21=t2[n]/t1[n];
		*/

		//see which prescale bin it is in
		if(n_real==2-1)
		for(unsigned int il=0; il<3; il++)
		  {
		    if(pt[n]> lhc_prescale_cut[il])
		      lhc_prescale[il].Add(w);
		  }

		n_real++;
	      }


	  }
      }
  cout<<"\r\033[K100% pre-processed."<<endl<<flush;


  //cout the total number of events in each bin
  cout<<"HT\\bin,A,B,C"<<endl;
  for(unsigned int k=1; k<=12; k++)
    {
      cout<<k;
      for(unsigned int ka=0; ka<3; ka++)
      {
	TChain chain_madgraph("Ntuple");
	stringstream sout;
	sout<<argv[1]<<alist[ka]<<"."<<k<<".root";
	chain_madgraph.Add(sout.str().c_str());
	
	//read the input file
	cout<<", "<<chain_madgraph.GetEntries();
      }
      cout<<endl;
    }

  cout<<"HT\\bin,A,B,C"<<endl;
  for(unsigned int k=1; k<=12; k++)
    {
      cout<<k;
      for(unsigned int ka=0; ka<3; ka++)
      {
	cout<<", "<<weights[k-1][ka];
      }
      cout<<endl;
    }



  //now compute all the prescale factors
  for(unsigned int il=0; il<3; il++)
    {
      //the number correspond to cxn in millibarns
      //first compute event rate
      //assume instantenous luminosity = 10^{-34} cm^{-2}
      // = 10 nanobarn^{-1} = 1e7 milibarn^{-1}
      double event_rate= 1e7 * lhc_prescale[il].nevt();

      cout<<"event rate factor for PT bin "<<lhc_prescale_cut[il]
	  <<" is: "<<lhc_prescale[il].nevt()<<endl;

      //make sure event_rate is positive
      assert(event_rate>0);

      //we now want to resctrict it to 200 Hz
      lhc_prescale_factor[il]= lhc_Hz/ event_rate;

      if(lhc_prescale_factor[il]>1.0)
	lhc_prescale_factor[il]=1.0;

      //make sure it is positive
      cout<<"Prescale factor for PT bin "<<lhc_prescale_cut[il]
	  <<" is: "<<lhc_prescale_factor[il]<<endl;
      assert(lhc_prescale_factor[il]>0);
      
      
    }


  //here we proceed to the actual training
  for(unsigned int k=1; k<=12; k++)
    for(unsigned int ka=0; ka<3; ka++)
      {
	TChain chain_madgraph("Ntuple");
	stringstream sout;
	sout<<argv[1]<<alist[ka]<<"."<<k<<".root";
	chain_madgraph.Add(sout.str().c_str());
	
	//read the input file
	chain_madgraph.SetBranchAddress("size", &size);
	chain_madgraph.SetBranchAddress("w", &w);
	chain_madgraph.SetBranchAddress("pt", pt);
	chain_madgraph.SetBranchAddress("eta", eta);
	chain_madgraph.SetBranchAddress("phi", phi);
	chain_madgraph.SetBranchAddress("m", m);
	chain_madgraph.SetBranchAddress("tau1", t1);
	chain_madgraph.SetBranchAddress("tau2", t2);
	chain_madgraph.SetBranchAddress("tau3", t3);
	
	double rescale=1.0/chain_madgraph.GetEntries();
	
	for(unsigned int i=0; i<chain_madgraph.GetEntries(); ++i)
	  {
	    if(i%10000==0)
	      cout<<"\r\033[K"<<floor((i/(chain_madgraph.GetEntries()*12.*3)  +
				       (3*(k-1)+ka)/(12.*3)
				       )*10000)*0.01
		  <<"% trained..."
		  <<flush;
	  
	    
	    chain_madgraph.GetEntry(i);
	    if(size<2) continue;

	    w*=rescale;
	    
	    int real_njet=0;
	    for(int n=0; n<size; ++n)
	      {
		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;
		if(pt[n] < min_pt) continue;
		if(m[n]<20) continue;
		real_njet++;
	      }
	    
	    if(real_njet!=n_exc) continue;
	    
	    //first find out prescale factor
	    int n_real=0;
	    double prescale_factor=1.0;
	    for(int n=0; n<size && n_real<n_exc && n_real<2; ++n)
	      {
		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;
		if(pt[n] < min_pt) continue;
		if(m[n]<20) continue;

		//if it is the second jet
		if(n_real==2-1)
		  for(unsigned int il=0; il<3; il++)
		    {
		      //find prescale bin
		      if(pt[n]>= lhc_prescale_cut[il] && 
			 (il>=2 || pt[n]< lhc_prescale_cut[il+1]))
			{
			  prescale_factor=lhc_prescale_factor[il];
			  //rescale the event
			  w*=prescale_factor;
			  break;
			}
		    }

		n_real++;
	      }

	    

	     n_real=0;
	    //limit to two jets
	    for(int n=0; n<size && n_real<n_exc && n_real<2; ++n)
	      {
		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;
		if(pt[n] < min_pt) continue;
		if(m[n]<20) continue;
		
		double logpt=log(pt[n]/min_pt);
		double mpt=-log10(m[n]/pt[n]);
		double t21=t2[n]/t1[n];
		
		dataset_mpt_njet[n_real]->Fill(make_list<double>(mpt,logpt),
					       w);
		dataset_tau21_njet[n_real]->Fill(make_list<double>(t21,logpt),
						 w);
		dataset_mpt_tau21_njet[n_real]->
		  Fill(make_list<double>(mpt,t21,logpt), w);

		//fill the statistics plots
		if(n_real<2)
		  {
		    for(int is=0; is<lhc_size; ++is)
		      {
			if(logpt< is+1 && logpt > is)
			  lhc_stat[is]->Fill(mpt, t21, w);
		      }
		  }


		//see contributions of different bins to files
		if(n_real==2-1)
		  {
		    if(pt[n]>50)
		      {
			error_contrib_50[k-1][ka] += w*w;
			nevt_contrib_50[k-1][ka] += w;
		      }
		    if(pt[n]>100)
		      {
			error_contrib_100[k-1][ka] += w*w;
			nevt_contrib_100[k-1][ka] += w;
		      }
		    if(pt[n]>500)
		      {
			error_contrib_500[k-1][ka] += w*w;
			nevt_contrib_500[k-1][ka] += w;
		      }
		  }
		n_real++;
	      }
	  }
	
      }
  cout<<"\r\033[K100% trained."<<endl<<flush;


  //statistics plots
  TFile file_stat("output_stat.root", "RECREATE");
  
  cout<<"Statistics plots stored in output_stat.root"<<endl;

  for(int is=0; is<lhc_size; ++is)
    {
      //scale it up and get the weights
      for(int i=0; i<=lhc_stat[is]->GetNbinsX(); ++i)
      for(int j=0; j<=lhc_stat[is]->GetNbinsY(); ++j)
	{
	  double temp_scale=0;
	  if(lhc_stat[is]->GetBinContent(i,j) > 0)
	    temp_scale=lhc_factor*
	      lhc_stat[is]->GetBinError(i,j)*lhc_stat[is]->GetBinError(i,j)/
	      (lhc_stat[is]->GetBinContent(i,j));
	  lhc_stat[is]->SetBinContent(i, j, temp_scale);
	}

      lhc_stat[is]->Write();
    }			



  file_stat.Close();


  //get error contribution
  for(unsigned int i=0; i<error_contrib_50.size(); ++i)
  for(unsigned int j=0; j<error_contrib_50[0].size(); ++j)
    {
      cout<<"For "<<i<<", "<<j<<endl;

      cout<<"50 GeV Nevt: "<<nevt_contrib_50[i][j]<<" +- "
	  <<sqrt(error_contrib_50[i][j])<<endl;

      cout<<"100 GeV Nevt: "<<nevt_contrib_100[i][j]<<" +- "
	  <<sqrt(error_contrib_100[i][j])<<endl;

      cout<<"500 GeV Nevt: "<<nevt_contrib_500[i][j]<<" +- "
	  <<sqrt(error_contrib_500[i][j])<<endl;

      cout<<endl;
    }

  //get PDF list
  vector<PDF> PDF_mpt_njet;
  vector<PDF> PDF_tau21_njet;
  vector<PDF> PDF_mpt_tau21_njet;
  
  //first turn-on Bias only in PDF computation
  DataSet::b_bias_only = true;
  
  //for the njet cases
  for(int n=0; n<dataset_mpt_njet.size() && n<n_exc && n<2; ++n)
    {
      PDF_mpt_njet.push_back(dataset_mpt_njet[n]->ComputePDF());
      PDF_tau21_njet.push_back(dataset_tau21_njet[n]->ComputePDF());      

      //if PDF smearing is on, smear the pdf first
      if(b_smear_PDF)
	{
	  //if second jet, the PDF is also computed
	  if(n==1)
	    continue;


	  //a negative scale_factor means no scaling
	  double scale_factor=-1.0;

	  if(b_smear_lhc)
	    {
	      scale_factor=lhc_factor;
	    }

	  DataSet dataset_new=
	    dataset_mpt_tau21_njet[n]->GenerateDataSet(scale_factor);

	  PDF_mpt_tau21_njet.push_back(dataset_new.ComputePDF());

	  //if first jet, add another PDF
	  if(n==0)
	    PDF_mpt_tau21_njet.push_back(dataset_new.ComputePDF());
	}

      else
	PDF_mpt_tau21_njet.push_back(dataset_mpt_tau21_njet[n]->ComputePDF());
    }

  
  //now free up some memory
  for(unsigned int n=0; n<dataset_mpt_njet.size(); ++n)
    {
      if(!dataset_mpt_njet[n])
	{
	  delete dataset_mpt_njet[n];
	  dataset_mpt_njet[n]=NULL;
	}	  
    }

  for(unsigned int n=0; n<dataset_tau21_njet.size(); ++n)
    {
      if(!dataset_tau21_njet[n])
	{
	  delete dataset_tau21_njet[n];
	  dataset_tau21_njet[n]=NULL;
	}	  
    }

  for(unsigned int n=0; n<dataset_mpt_tau21_njet.size(); ++n)
    {
      if(!dataset_mpt_tau21_njet[n])
	{
	  delete dataset_mpt_tau21_njet[n];
	  dataset_mpt_tau21_njet[n]=NULL;
	}
    }


  //now from the exclusive 3 jet sample, we will proceed to get actual
  //data-driven background in the 4 jet sample

  //first declare dataset
  //for 2D distributions we want later


  //create generators for MC toy integration
  vector<PDF::MC_Generator> MC_mpt_tau21_njet;
  
  for(int n=0; n< PDF_mpt_njet.size(); ++n)
    {
      MC_mpt_tau21_njet.push_back(PDF_mpt_tau21_njet[n].Generator
				  (make_list<PDF::MC_Flag>
				   (PDF::OUTPUT, PDF::OUTPUT, PDF::INPUT)
				   ));
    }


  //make sure n_exc=3
  //we are doing 4 jet samples
  assert(n_exc==3);

  //declare 2D distributions to study
  DataSet t21_j12(make_list<double>(0, 0),
		  make_list<double>(1.0, 1.0),
		  make_list<double>(0.02, 0.02));
  
  DataSet t21_j34(make_list<double>(0, 0),
		  make_list<double>(1.0, 1.0),
		  make_list<double>(0.02, 0.02));

  DataSet mpt_j12(make_list<double>(0, 0),
		  make_list<double>(1.5, 1.5),
		  make_list<double>(0.01, 0.01));

  DataSet mpt_j34(make_list<double>(0, 0),
		  make_list<double>(1.5, 1.5),
		  make_list<double>(0.01, 0.01));


  DataSet t21_j12_MC(make_list<double>(0, 0),
		     make_list<double>(1.0, 1.0),
		     make_list<double>(0.02, 0.02));

  DataSet t21_j34_MC(make_list<double>(0, 0),
		     make_list<double>(1.0, 1.0),
		     make_list<double>(0.02, 0.02));
  
  DataSet mpt_j12_MC(make_list<double>(0, 0),
		     make_list<double>(1.5, 1.5),
		     make_list<double>(0.01, 0.01));
  
  DataSet mpt_j34_MC(make_list<double>(0, 0),
		     make_list<double>(1.5, 1.5),
		     make_list<double>(0.01, 0.01));


  //declare a bunch of counters
  //first set the cuts
  vector<double> T21_cut;
  vector<double> MJ_cut;
  vector<Counter> nevt_real;
  vector<Counter> nevt_MC;


  //plots for the corresponding MJ and T21 distributions
  vector<double> plot_MJ_cut;
  vector<double> plot_T21_cut;

  vector<TH1D*> plot_MJ_real;
  vector<TH1D*> plot_T21_real;

  vector<TH1D*> plot_MJ_MC;
  vector<TH1D*> plot_T21_MC;

  vector<TH1D*> plot_MJ_unbiased_MC;
  vector<TH1D*> plot_T21_unbiased_MC;



  //now do T21 cut from 0.2 to 1.0 in 0.1 increments (<- no cut)
  //jet mass cut from 0 to 500 GeV in 50 GeV increments
  bool b_first=true;
  
  for(double T21cut=0.2; T21cut<1.01; T21cut+=0.1)
    {
      //grab the T21 cut
      plot_T21_cut.push_back(T21cut);
      stringstream sout, sout_tit;
      
      //grab the histograms
      sout<<"plot_MJ_"<<T21cut<<"_real";
      sout_tit<<"M_{J} plot after T_{21} > "<<T21cut<<" (real)";
      TH1D* temp_plot_real= new TH1D(sout.str().c_str(),
				     sout_tit.str().c_str(),
				     200, 0, 2000);

      sout.str("");
      sout_tit.str("");
      sout<<"plot_MJ_"<<T21cut;
      sout_tit<<"M_{J} plot after T_{21} > "<<T21cut;
      TH1D* temp_plot= new TH1D(sout.str().c_str(),
				sout_tit.str().c_str(),
				200, 0, 2000);


      sout<<"_b";
      sout_tit<<" (unbiased)";
      TH1D* temp_plot_b=new TH1D(sout.str().c_str(),
				 sout_tit.str().c_str(),
				 200, 0, 2000);
			   

      temp_plot_real->Sumw2();
      temp_plot->Sumw2();
      temp_plot_b->Sumw2();

      plot_MJ_real.push_back(temp_plot_real);
      plot_MJ_MC.push_back(temp_plot);
      plot_MJ_unbiased_MC.push_back(temp_plot_b);

      for(double MJcut=100; MJcut<1001; MJcut+=50.)
	{
	  T21_cut.push_back(T21cut);
	  MJ_cut.push_back(MJcut);
	  nevt_real.push_back(Counter());
	  nevt_MC.push_back(Counter());	

	  //grab the MJ cut
	  if(b_first)
	    {
	      plot_MJ_cut.push_back(MJcut);

	      //grab the histograms
	      sout.str("");
	      sout_tit.str("");
	      sout<<"plot_T21_"<<MJcut<<"_real";
	      sout_tit<<"T_{21} plot after M_{J} > "<<MJcut<<" (real)";
	      TH1D* temp_plott_real= new TH1D(sout.str().c_str(),
					      sout_tit.str().c_str(),
					      20, 0, 1.0);
	      
	      sout.str("");
	      sout_tit.str("");
	      sout<<"plot_T21_"<<MJcut;
	      sout_tit<<"T_{21} plot after M_{J} > "<<MJcut;
	      TH1D* temp_plott= new TH1D(sout.str().c_str(),
					sout_tit.str().c_str(),
					20, 0, 1.0);
	      
	      sout<<"_b";
	      sout_tit<<" (unbiased)";
	      TH1D* temp_plott_b=new TH1D(sout.str().c_str(),
					  sout_tit.str().c_str(),
					  20, 0, 1.0);

	      temp_plott_real->Sumw2();
	      temp_plott->Sumw2();
	      temp_plott_b->Sumw2();

	      plot_T21_real.push_back(temp_plott_real);
	      plot_T21_MC.push_back(temp_plott);
	      plot_T21_unbiased_MC.push_back(temp_plott_b);
	    }
	  
	}
      
      b_first=false;
    }

  cout<<"starting toy MC integration..."<<endl;

  for(unsigned int k=1; k<=12; k++)
    for(unsigned int ka=1; ka<3; ka++)
      {
	TChain chain_madgraph("Ntuple");
	stringstream sout;
	sout<<argv[1]<<alist[ka]<<"."<<k<<".root";
	chain_madgraph.Add(sout.str().c_str());
	
	//read the input file
	chain_madgraph.SetBranchAddress("size", &size);
	chain_madgraph.SetBranchAddress("w", &w);
	chain_madgraph.SetBranchAddress("pt", pt);
	chain_madgraph.SetBranchAddress("eta", eta);
	chain_madgraph.SetBranchAddress("phi", phi);
	chain_madgraph.SetBranchAddress("m", m);
	chain_madgraph.SetBranchAddress("tau1", t1);
	chain_madgraph.SetBranchAddress("tau2", t2);
	chain_madgraph.SetBranchAddress("tau3", t3);
	
	double rescale=1.0/chain_madgraph.GetEntries();
	
	for(unsigned int i=0; i<chain_madgraph.GetEntries(); ++i)
	  {
	    if(i%1000==0)
	      cout<<"\r\033[K"<<floor((i/(chain_madgraph.GetEntries()*12.*3)  +
				       (3*(k-1)+ka)/(12.*3)
				       )*10000)*0.01
		  <<"% processed..."
		  <<flush;

	    chain_madgraph.GetEntry(i);
	    if(size<2) continue;
	    
	    w*=rescale;
	    
	    //if doing lhc scaling, rescale the weight
	    if(b_smear_lhc)
	      w*=lhc_factor;
	  
	    int real_njet=0;

	    double logpt_real[4];
	    double mpt_real[4];
	    double t21_real[4];

	    double T21_evt=1.0;
	    double MJ_evt=0;

	    for(int n=0; n<size; ++n)
	      {
		if(t1[n] < 1e-5) continue;
		if(t2[n] < 1e-5) continue;

		//actual pt cut is 20 GeV higher
		if(pt[n] < min_pt + 50.0) continue;
		if(m[n]<20) continue;

		//do only first two jets
		if(real_njet < 4)
		  {
		    logpt_real[real_njet]=log(pt[n]/min_pt);
		    mpt_real[real_njet]=-log10(m[n]/pt[n]);
		    t21_real[real_njet]=t2[n]/t1[n];
		    
		    
		    if(real_njet < 2)
		      {
			T21_evt*= sqrt(t21_real[real_njet]);
			MJ_evt+= m[n];
		      }
		    
		  }

		real_njet++;
	      }

	    //we want the signal region to be
	    //n_exc + 1
	    if(real_njet < n_exc+1) continue;
	    
	    //now populate the DataSet
	    t21_j12.Fill(make_list<double>(t21_real[0], t21_real[1]), w);
	    t21_j34.Fill(make_list<double>(t21_real[2], t21_real[3]), w);

	    mpt_j12.Fill(make_list<double>(mpt_real[0], mpt_real[1]), w);
	    mpt_j34.Fill(make_list<double>(mpt_real[2], mpt_real[3]), w);
	    
	    //fill the counters
	    for(unsigned int n=0; n<T21_cut.size(); n++)
	      {
		if(T21_evt < T21_cut[n] &&
		   MJ_evt > MJ_cut[n])
		  nevt_real[n].Add(w);
	      }

	    //fill the plots
	    for(unsigned int n=0; n<plot_MJ_cut.size(); n++)
	      if(MJ_evt > plot_MJ_cut[n])
		plot_T21_real[n]->Fill(T21_evt, w);

	    for(unsigned int n=0; n<plot_T21_cut.size(); n++)
	      if(T21_evt < plot_T21_cut[n])
		plot_MJ_real[n]->Fill(MJ_evt, w);
	    

	    //if not MC integrate, skip this part
	    if(!do_MC)
	      continue;

	    //we have an event now!
	    //now the first three jets, we have templates
	    //so create a vector for them
	    vector<vector<double> > MC_pt;
	    vector<vector<double> > MC_m;
	    vector<vector<double> > MC_mpt;
	    vector<vector<double> > MC_tau21;
	    vector<vector<double> > MC_weight;
	    vector<vector<double> > MC_weight_bias;
	    vector<vector<double> > MC_err;
	    
	    int nGenerate=1000;

	    //now populate the MC data
	    //limit to two jets
	    for(int n=0; n<n_exc+1 && n<size && n<2; ++n)
	      {
		//we use 3rd jet template for 4th jet
		int actual_n=n;
		if(n==n_exc)
		  actual_n=n_exc-1;

		vector<PDF::MC_Event> MC_mpt_tau21=
		  MC_mpt_tau21_njet[actual_n].GenerateFast
		  (make_list<double>(log(pt[n]/min_pt)), nGenerate);

		MC_m.push_back(vector<double>(MC_mpt_tau21.size()));
		MC_mpt.push_back(vector<double>(MC_mpt_tau21.size()));
		MC_tau21.push_back(vector<double>(MC_mpt_tau21.size()));
		MC_weight.push_back(vector<double>(MC_mpt_tau21.size()));
		MC_weight_bias.push_back(vector<double>(MC_mpt_tau21.size()));
		MC_err.push_back(vector<double>(MC_mpt_tau21.size()));

		for(unsigned int nMC=0; nMC<MC_mpt_tau21.size(); ++nMC)
		  {
		    double t_mass=exp10(-MC_mpt_tau21[nMC].value[0])*pt[n];
		    double t_tau21=MC_mpt_tau21[nMC].value[1];
		   
		    //if mass is < 20, set weights to zero
		    if(t_mass < 20.)
		      {
			MC_weight.back().push_back(0);
			MC_weight_bias.back().push_back(0);
			MC_err.back().push_back(0);
		      }

		    else
		      {
			MC_weight.back().push_back(MC_mpt_tau21[nMC].weight);
			MC_weight_bias.back().push_back
			  (MC_mpt_tau21[nMC].weight-MC_mpt_tau21[nMC].err);
			MC_err.back().push_back(MC_mpt_tau21[nMC].err);
		      }

 
		    MC_m.back().push_back(t_mass);
		    MC_mpt.back().push_back(MC_mpt_tau21[nMC].value[0]);
		    MC_tau21.back().push_back(t_tau21);
		  }
	      }


	    //first reweight the events to make the sum zero
	    double weight12=0;
	    double weight12_bias=0;

	    for(unsigned int n=0; n<MC_weight[0].size(); ++n)
	      {
		weight12 += MC_weight[0][n]*MC_weight[1][n];
		weight12_bias += MC_weight_bias[0][n]*MC_weight_bias[1][n];
	      }

	    assert(weight12!=0);
	    assert(weight12_bias!=0);

	    weight12=1.0/weight12;
	    weight12_bias=1.0/weight12_bias;
	    
	    //now we have a bunch of toy MC events
	    //populate the DataSet using MC events
	    
	    //first we create temporary vector to store error
	    vector<double> nevt_MC_temp(nevt_MC.size());
	    vector<double> nevt_err_temp(nevt_MC.size());
	    
	    for(unsigned int n=0; n<MC_weight[0].size(); ++n)
	      {
		//now populate the DataSet
		t21_j12_MC.Fill(make_list<double>
				(MC_tau21[0][n], MC_tau21[1][n]),
				w*weight12*
				MC_weight[0][n]*MC_weight[1][n]);

		/*
		t21_j34_MC.Fill(make_list<double>
				(MC_tau21[2][n], MC_tau21[3][n]),
				w*weight34*
				MC_weight[2][n]*MC_weight[3][n]);
		*/

		mpt_j12_MC.Fill(make_list<double>
				(MC_mpt[0][n], MC_mpt[1][n]), 
				w*weight12*
				MC_weight[0][n]*MC_weight[1][n]);

		/*
		mpt_j34_MC.Fill(make_list<double>
				(MC_mpt[2][n], MC_mpt[3][n]), 
				w*weight34*
				MC_weight[2][n]*MC_weight[3][n]);
		*/

		double T21_evt=sqrt(MC_tau21[0][n]*MC_tau21[1][n]);
				   

		double MJ_evt=MC_m[0][n]+MC_m[1][n];//+MC_m[2][n];
		double weight=MC_weight[0][n]*MC_weight[1][n]*weight12;

		double err= MC_weight_bias[0][n]*MC_weight_bias[1][n]*
		  weight12_bias;

		
		for(unsigned int ncut=0; ncut<T21_cut.size(); ncut++)
		  {
		    if(T21_evt < T21_cut[ncut] &&
		       MJ_evt > MJ_cut[ncut])
		      {
			nevt_MC_temp[ncut]+=weight;
			nevt_err_temp[ncut]+=err;

			//add the erros to nevt_mC
			if(b_smear_lhc)
			  nevt_MC[ncut].Add_tot_n_lhc(w, weight);

			else
			  nevt_MC[ncut].Add_tot_n(w*weight);
		       
		      }
		  }

		//populate all the plots
		for(unsigned int n=0; n<plot_MJ_cut.size(); n++)
		  if(MJ_evt > plot_MJ_cut[n])
		    {
		      plot_T21_MC[n]->Fill(T21_evt, w*weight);
		      plot_T21_unbiased_MC[n]->Fill(T21_evt, w*err);
		    }

		for(unsigned int n=0; n<plot_T21_cut.size(); n++)
		  if(T21_evt < plot_T21_cut[n])
		    {
		      plot_MJ_MC[n]->Fill(MJ_evt, w*weight);
		      plot_MJ_unbiased_MC[n]->Fill(MJ_evt, w*err);
		    }

	      } // loop over all MC events

	    /*
	    cout<<"total weight: "<<nevt_MC_temp[50]
		<<","<<nevt_err_temp[50]<<endl;
	    */

	    //now fill counter
	    for(unsigned int ncut=0; ncut<nevt_MC.size(); ++ncut)
	      {
		if(b_smear_lhc)
		  nevt_MC[ncut].Add_lhc(w, nevt_MC_temp[ncut],
					w*(nevt_MC_temp[ncut]-nevt_err_temp[ncut]));
		
		else
		  nevt_MC[ncut].Add(w*nevt_MC_temp[ncut],
				    w*(nevt_MC_temp[ncut]-nevt_err_temp[ncut]));
	      }
	  }



      }
  cout<<"\r\033[Kfinished all processing"<<endl<<flush;
  

  cout<<"Writing graphs..."<<endl;

  //store some of the templates
  TFile out(argv[2], "RECREATE");


  if(do_MC)
    {
      plot2D(t21_j12, "t21_j12").Write();
      plot2D(t21_j12_MC, "t21_j12_MC").Write();
      //plot2D(t21_j34, "t21_j34").Write();
      //plot2D(t21_j34_MC, "t21_j34_MC").Write();
      
      plot2D(mpt_j12, "mpt_j12").Write();
      plot2D(mpt_j12_MC, "mpt_j12_MC").Write();
      //  plot2D(mpt_j34, "mpt_j34").Write();
      //  plot2D(mpt_j34_MC, "mpt_j34_MC").Write();
    }

  if(do_MC)
    out.mkdir("templates")->cd();
  map<string, Vec> graph_list;
  
  for(int i=60; i<=3000; i+=10)
    {
      stringstream sout;
      
      //do the njet case
      if(do_MC)
      for(int n=0; n<2 && n<n_exc; ++n)
	{
	  sout.str("");
	  sout<<"t21_"<<n<<"jet_"<<i;
	  graph_list[sout.str()].fill_data(PDF_tau21_njet[n], i);
	  graph_list[sout.str()].get_graph(sout.str()).Write();
	  
	  sout.str("");
	  sout<<"mpt_"<<n<<"jet_"<<i;
	  graph_list[sout.str()].fill_data(PDF_mpt_njet[n], i);
	  graph_list[sout.str()].get_graph(sout.str(), true).Write();	  
	}
      
    }
    
  
  out.mkdir("T21_plots")->cd();
  for(unsigned int n=0; n<plot_MJ_cut.size(); n++)
    {
      if(do_MC)
	{
	  plot_T21_MC[n]->Write();
	  plot_T21_unbiased_MC[n]->Write();
	}
      plot_T21_real[n]->Write();
    }

  out.mkdir("MJ_plots")->cd();
  for(unsigned int n=0; n<plot_T21_cut.size(); n++)
    {
      if(do_MC)
	{
	  plot_MJ_MC[n]->Write();
	  plot_MJ_unbiased_MC[n]->Write();
	}
      plot_MJ_real[n]->Write();
    }

  out.Close();

  ofstream fout("cut_table.txt");

  fout<<"T21_cut, MJ_cut, nevt_real, nevt_real_err, nevt_MC, nevt_MC_bias, stat_err, toy_MC_err"<<endl;

  double output_factor=1e9;

  if(b_smear_lhc)
    output_factor=1.0;

  for(unsigned int n=0; n<nevt_real.size(); ++n)
    {
      fout<<T21_cut[n]<<", "<<MJ_cut[n]<<", "
	  <<nevt_real[n].nevt()*output_factor<<", "
	  <<nevt_real[n].err()*output_factor<<", "
	  <<nevt_MC[n].nevt()*output_factor<<", "
	  <<nevt_MC[n].bias()*output_factor<<", "
	  <<nevt_MC[n].var()*output_factor<<", "
	  <<nevt_MC[n].tot_err()
	  <<endl;
    }

  fout.close();

  cout<<"Program ends."<<endl;
  return 1;
}
  
